package fem2.strategies;

import fem2.Debugger;
import fem2.Model;
import fem2.enu.EchoLevelType;
import fem2.enu.SolverType;
import fem2.jlinalg.Solver;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IVector;
import inf.jlinalg.SolveFailedException;

import java.util.Arrays;

import math2.MathUtilities;

/**
 * Arc length solution strategy
 * 
 * @author hbui
 * 
 */
public abstract class ArcLengthSolutionStrategy extends Strategy {

	private ArrayVector r = new ArrayVector();
	private double nr;
	private ArrayVector uHat = null; // play the role of un
	private ArrayVector uHatPredictor = null; // play the role of un_predictor
	private ArrayVector du_lambda = null;
	private ArrayVector dfdu = null;
	private double lambda;
	private double lambdaOld; // lambdan
	private double lambdaPredictor; // lambdan_predictor
	private double constraintValue = 0.0;

	public ArcLengthSolutionStrategy(Model m) {
		super(m);
		// setSolverType(SolverType.UMFPACK_SOLVER);
		setSolverType(SolverType.MKL_SOLVER);
	}

	@Override
	public void setSolverType(SolverType type) {
		if ((type != SolverType.UMFPACK_SOLVER) && (type != SolverType.MKL_SOLVER)) {
			throw new Error(
					"For efficiency reason (stiffness matrix is reused for multiple solutions). This strategy does not allow using solver other than "
							+ SolverType.UMFPACK_SOLVER + " or " + SolverType.MKL_SOLVER);
		}
		super.setSolverType(type);
	}

	@Override
	public double getCurrentTimeStep() {
		return lambda;
	}

	@Override
	public double[] getCurrentUHat() {
		return uHat.getValues();
	}

	@Override
	public double[] getCurrentUdHat() {
		throw new Error("this method is not available for this strategy");
	}

	@Override
	public double[] getCurrentUddHat() {
		throw new Error("this method is not available for this strategy");
	}

	@Override
	public void started(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * initialize model (enumeration)
		 */
		super.started(leftHandSideSolver, rightHandSide, u);
		int n = model.getSize();

		/*
		 * initialize solver
		 */
		initializeSolver(leftHandSideSolver);

		/*
		 * initialize auxiliary vector
		 */
		r.setSize(n);
		u.setSize(n);
		rightHandSide.setSize(n);
		uHat = new ArrayVector(n);
		uHatPredictor = new ArrayVector(n);
		du_lambda = new ArrayVector(n);
		dfdu = new ArrayVector(n);

		/*
		 * compute the external force vector. In the arc-length solution
		 * strategy, external force vector is not allowed to be changed since
		 * the method automatically track the response of the system. Therefore
		 * a consistent external force vector during time step iteration is
		 * required. By default, external force vector is sampled at t = 0.0.
		 */
		model.assembleR(r.getValues(), 0.0);
		nr = BLAM.norm(BLAM.NORM_TWO, r);

		/*
		 * initialize model at time step 0.0
		 */
		model.setX(0.0, u.getValues());
		model.InitializeSolutionStep();

		/*
		 * assemble initial tangent stiffness and right hand side vector
		 */
		leftHandSideSolver.initialize();
		model.assembleKt(leftHandSideSolver.getA());
		// Arrays.fill(rightHandSide.getValues(), 0.0);
		// model.assembleRi(rightHandSide.getValues());
		// BLAM.scale(-1.0, rightHandSide);
		// BLAM.add(lambda, r, rightHandSide);
	}

	@Override
	public void timeStepStarted(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * perform predictor using secant path criterion. (i.e Souza Neto)
		 */
		BLAM.copy(r, rightHandSide);
		try {
			leftHandSideSolver.solve(rightHandSide.getValues());
		} catch (SolveFailedException e) {
			throw new Error("Solve failed");
		}

		double s = getNextSteppingParameter();
		double s0 = Math.sqrt(BLAM.dot(rightHandSide, rightHandSide) + 1);
		double psi = getScalingParameter();
		double forward_criteria = BLAM.dot(rightHandSide, uHat)
				- BLAM.dot(rightHandSide, uHatPredictor) + psi * (lambdaOld - lambdaPredictor);
		if (forward_criteria < 0) {
			s0 = -s0;
		}

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("forward_criteria = ", forward_criteria);
			Debugger.watch("s = ", s);
			Debugger.watch("s0 = ", s0);
			Debugger.watch("delta_u_lambda = ", rightHandSide);
		}

		lambda += s / s0;
		BLAM.add(s / s0, rightHandSide, u);

		/*
		 * save predictor state
		 */
		lambdaPredictor = lambda;
		BLAM.copy(u, uHatPredictor);

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("u_predictor = ", u);
			Debugger.watch("lambda_predictor = ", lambda);
		}

		/*
		 * initialize model at time step start
		 */
		model.setX(0.0, u.getValues());
		model.InitializeSolutionStep();

		/*
		 * compute initial rightHandSide
		 */
		Arrays.fill(rightHandSide.getValues(), 0.0);
		model.assembleRi(rightHandSide.getValues());
		BLAM.scale(-1.0, rightHandSide);
		BLAM.add(lambda, r, rightHandSide);

		/*
		 * compute initial value of constraint
		 */
		constraintValue = computeConstraintValue(u, uHat, lambda, lambdaOld);
	}

	@Override
	public void iterationStarted(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		model.InitalizeNonlinearIteration();

		/*
		 * zero out and assemble tangent stiffness
		 */
		leftHandSideSolver.initialize();
		model.assembleKt(leftHandSideSolver.getA());

		/*
		 * compute delta_u_lambda
		 */
		BLAM.copy(r, du_lambda);
		try {
			leftHandSideSolver.solve(du_lambda.getValues());
		} catch (SolveFailedException e) {
			throw new Error("Solve failed");
		}

		/*
		 * compute delta_u_r
		 */
		try {
			leftHandSideSolver.solve(rightHandSide.getValues());
		} catch (SolveFailedException e) {
			throw new Error("Solve failed");
		}

	}

	@Override
	public boolean iterationFinished(Solver leftHandSideSolver, ArrayVector rightHandSide,
			ArrayVector u) {
		/*
		 * compute delta_lambda
		 */
		computeGradientOfConstraintWrtU(u, uHat, lambda, lambdaOld, dfdu);
		double dfdlambda = computeDerivativeOfConstraintWrtLambda(u, uHat, lambda, lambdaOld);
		double dlambda = -(constraintValue + BLAM.dot(dfdu, rightHandSide))
				/ (BLAM.dot(dfdu, du_lambda) + dfdlambda);

		/*
		 * update u
		 */
		BLAM.add(dlambda, du_lambda, u);
		BLAM.add(1.0, rightHandSide, u);

		/*
		 * update lambda
		 */
		lambda += dlambda;

		/*
		 * update model/material
		 */
		model.setX(0.0, u.getValues());
		model.FinalizeNonlinearIteration(); // update material

		/*
		 * compute rightHandSide
		 */
		Arrays.fill(rightHandSide.getValues(), 0.0);
		model.assembleRi(rightHandSide.getValues());
		BLAM.scale(-1.0, rightHandSide);
		BLAM.add(lambda, r, rightHandSide);

		/*
		 * compute constraint value
		 */
		constraintValue = computeConstraintValue(u, uHat, lambda, lambdaOld);

		/*
		 * check convergence criteria
		 */
		double res;
		if (nr > MathUtilities.tol12) {
			res = Math.sqrt(Math.pow(BLAM.norm(BLAM.NORM_TWO, rightHandSide) / nr, 2)
					+ Math.pow(constraintValue, 2));
		} else {
			res = Math.sqrt(Math.pow(BLAM.norm(BLAM.NORM_TWO, rightHandSide), 2)
					+ Math.pow(constraintValue, 2));
		}

		boolean conv_criteria = res < MathUtilities.tol6;

		if (conv_criteria) {
			if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
				Debugger.watch("u_corrector = ", u);
				Debugger.watch("lambda_corrector = ", lambda);
			}
		}

		return conv_criteria;
	}

	@Override
	public void timeStepFinished(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		model.FinalizeSolutionStep();
		/*
		 * save solution vector
		 */
		BLAM.copy(u, uHat);
		/*
		 * save lambda
		 */
		lambdaOld = lambda;
	}

	/**
	 * get stepping parameter of the arc length method
	 * 
	 * @return
	 */
	protected abstract double getNextSteppingParameter();

	/**
	 * get current stepping parameter of the arc length method
	 * 
	 * @return
	 */
	protected abstract double getCurrentSteppingParameter();

	/**
	 * get scaling parameter between solution vector and control parameter
	 * lambda
	 * 
	 * @return
	 */
	protected abstract double getScalingParameter();

	/**
	 * compute f(u, lambda)
	 * 
	 * @param u
	 * @param u_old
	 * @param lambda
	 * @param lambda_old
	 * @return
	 */
	protected abstract double computeConstraintValue(IVector u, IVector u_old, double lambda,
			double lambda_old);

	/**
	 * compute df/du
	 * 
	 * @param u
	 * @param u_old
	 * @param lambda
	 * @param lambda_old
	 * @return
	 */
	protected abstract void computeGradientOfConstraintWrtU(IVector u, IVector u_old,
			double lambda, double lambda_old, IVector dfdu);

	/**
	 * compute df/dlambda
	 * 
	 * @param u
	 * @param u_old
	 * @param lambda
	 * @param lambda_old
	 * @return
	 */
	protected abstract double computeDerivativeOfConstraintWrtLambda(IVector u, IVector u_old,
			double lambda, double lambda_old);

}
